import cv2
import os


def clip_video_by_sec(source_video, target_video, start_point, end_point, save_fourcc_format='mp4v'):
    cap = cv2.VideoCapture(source_video)
    if cap.isOpened():
        f_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
        f_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
        f_shape = (int(f_width), int(f_height))
        fps = round(cap.get(cv2.CAP_PROP_FPS))
        print("fps: ", fps)
        frame_number = cap.get(cv2.CAP_PROP_FRAME_COUNT)
        duration = frame_number / fps
        start_frame = int(fps * start_point)
        end_frame = int(fps * end_point)
        four_cc = cv2.VideoWriter_fourcc(*save_fourcc_format)
        video_writer = cv2.VideoWriter(target_video, four_cc, fps, f_shape)
        count = 0
        while True:
            _, frame = cap.read()
            if frame is None:
                break
            if count <= int(start_frame):
                print('Counting number ' + str(count + 1) + ' frequency')
            elif int(start_frame) < count <= int(end_frame):
                print('Writing number ' + str(count + 1) + ' frequency')
                video_writer.write(frame)
            else:
                break
            count += 1
        cap.release()
        video_writer.release()
        print("FPS: ", fps)
        print("f_shape: ", f_width, f_height)
        print(str(start_point // 60) + ":" + str(start_point % 60) + " - " + str(end_point // 60) + ":" + str(
            end_point % 60))
        print("Frame Number: ", frame_number)
        print("Duration: ", duration)
    else:
        print('video opened failed !')


def clip_video_by_frame(source_video, target_video, fps, start_frame, end_frame, save_fourcc_format='mp4v'):
    cap = cv2.VideoCapture(source_video)
    if cap.isOpened():
        f_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
        f_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
        f_shape = (int(f_width), int(f_height))
        four_cc = cv2.VideoWriter_fourcc(*save_fourcc_format)
        video_writer = cv2.VideoWriter(target_video, four_cc, fps, f_shape)
        count = 0
        while cap.isOpened():
            _, frame = cap.read()
            if count <= int(start_frame):
                print('Counting number ' + str(count + 1) + ' frequency')
            elif int(start_frame) < count <= int(end_frame):
                print('Writing number ' + str(count + 1) + ' frequency')
                video_writer.write(frame)
            else:
                break
            count += 1
        cap.release()
        video_writer.release()
        print("FPS: ", fps)
        print("f_shape: ", f_width, f_height)
        print(str((start_frame / fps) // 60) + ":" + str((start_frame / fps) % 60) + " - " + str(
            (end_frame / fps) // 60) + ":" + str((end_frame / fps) % 60))
        print("file path : ", target_video)
    else:
        print('video opened failed !')


if __name__ == "__main__":
    root_path = r'/home/kaijia/algo-env/Utils/sample/256A'
    video_name = r"256A-2024-03-27_09-13-08.mp4"
    file_name, _ = os.path.splitext(video_name)
    source_file = os.path.join(root_path, video_name)
    storage_path = os.path.join(root_path, "video_fragment")
    if not os.path.exists(storage_path):
        os.makedirs(storage_path)
    ########## 以秒数为准剪辑视频 ##########
    start_time = 17 * 60 + 20
    end_time = 18 * 60 + 10
    target_name = str(start_time) + "_" + str(end_time)
    target_file = os.path.join(storage_path, file_name + "_" + "sec_" + target_name + ".mp4")
    print("Video file storages in ", target_file)
    clip_video_by_sec(source_file, target_file, start_time, end_time)
    ########## 以帧数为准剪辑视频 ##########
    # start_f = 984
    # end_f = 1373
    # target_name = str(start_f) + "_" + str(end_f)
    # target_file = os.path.join(storage_path, file_name + "_" + "frame_" + target_name + ".mp4")
    # clip_video_by_frame(source_file, target_file, 25, start_f, end_f)
